import os
import sys
sys.path.append(os.path.abspath(os.path.dirname(__file__) + '/' + '../..'))

import matplotlib as mpl

mpl.use('pgf')
pgf_with_latex = {                     
    "pgf.texsystem": "xelatex",        
    "pgf.rcfonts": False,
    "text.usetex": True,                
    "font.family": "Times New Roman",
    "pgf.preamble": [
        r"\usepackage{fontspec}",    
        r"\setmainfont{Times New Roman}",        
        r"\usepackage{unicode-math}",
        r"\setmathfont{XITS Math}"
        ]
    }    
mpl.rcParams.update(pgf_with_latex)

import numpy as np
import matplotlib.pyplot as plt

from python.tools import (
    clean_folder
)

from python.numerical.functions import (
    calculate_bias_coefs,
    calculate_composite_bias_coefs_noisy
)

################
## Parameters ##
################

output_folder = './numerical/output/measurement_error'

params_DGP = {'rho': 0.80,
              'sigma_eps': 1.0} # Normalization
params_DGP['sigma_omega'] = 3.0 * params_DGP['sigma_eps']
L = 10 # Number of bias coefficients to plot
sigma_m_L = 1.0 # Measurement error in "low" scenario
sigma_m_H = 2.0 # Measurement error in "high" scenario

params_noisy_diagnostic = {'rho_perceived': params_DGP['rho'],
                           'sigma_eps_perceived': params_DGP['sigma_eps'],
                           'sigma_omega_perceived': params_DGP['sigma_omega'],
                           'theta': 0.60}

###############
## Get plots ##
###############

clean_folder(output_folder)

# Calculate bias coefficients
df = calculate_bias_coefs(params_DGP, params_noisy_diagnostic, L)
df['composite_meas_error_low'] = calculate_composite_bias_coefs_noisy(params_DGP, params_noisy_diagnostic, L, sigma_m = sigma_m_L)
df['composite_meas_error_high'] = calculate_composite_bias_coefs_noisy(params_DGP, params_noisy_diagnostic, L, sigma_m = sigma_m_H)

# Save output
df.to_csv('{}/meas_error.csv'.format(output_folder), index = False)

# Get graph
fig, ax = plt.subplots(figsize=(4.5, 3.0))
plt.plot(df['L'], df['composite_meas_error_high'], 'k:v', label = 'High measurement error')
plt.plot(df['L'], df['composite_meas_error_low'], 'k--^', label = 'Low measurement error')
plt.plot(df['L'], df['composite'], 'k-o', label = 'No measurement error')
plt.xlabel('Lag', fontsize = 12)
plt.ylabel('Composite Bias Coefficients', fontsize = 12)
plt.axhline(0, color = 'k', linestyle = '--', alpha = 0.50)
plt.legend(frameon = False, fontsize = 12)
plt.tight_layout()
plt.savefig('{}/meas_error.png'.format(output_folder), dpi = 600)
plt.savefig('{}/meas_error.pgf'.format(output_folder), dpi = 600)